using System;
using System.Collections;
using System.Reflection;
using NBody;
using NBody.Analysis;
using OptionParser;

[assembly: AssemblyTitle ("NBodyBullock")]
[assembly: AssemblyVersion ("1.0.*")]
[assembly: AssemblyDescription ("Compute the specific angular momentum.")]
[assembly: AssemblyCopyright ("2005 Joseph D. MacMillan")]

public class BullockOptions : InputOutputGetOptions
{
    [Option ("-n", "Number of particles per radial bin.")]
    public int NumPerBin;
    
    [Option ("-c", "Adjust for centre of mass")]
    public bool CentreOfMass;
    
    [Option ("-a", "Check alignment of angular momentum in halo")]
    public bool Alignment;
    
    [Option ("-s", "Use j^2 rather than jz")]
    public bool Squared;
    
    [Option ("-b", "Bin in spherical shells instead")]
    public bool Bin;
    
    [Option ("-r", "output relative values")]
    public bool Relative;
        
    public BullockOptions()
    {
        NumPerBin = 1500;
    }
}

public class NBodyBullock
{
    static bool use_squared;
    
    public static void Main(string[] args)
    {
        BullockOptions opts = new BullockOptions();
        opts.ProcessArgs(args, ".bull");
        use_squared = opts.Squared;
                        
        NBodySystem s = NBodySystem.Read(opts.InFile);
        
        if (opts.CentreOfMass)
        {
            s.AdjustForCentreOfMass(1e-3);
            NBody.Cosmology.Model cosmo = NBody.Cosmology.Model.SCDM;
            double z = cosmo.Redshift(s.Time);
            double rho_c = cosmo.Density(z);
            double r_200 = VirialRadius.CalculateR200(s, rho_c);
            s.ExtractSphere(r_200);
        }
        
        double jx, jy, jz;
        Bullock.CalculateTotalJVector(s, out jx, out jy, out jz);
        double jtot = Math.Sqrt(jx*jx + jy*jy + jz*jz);
        Console.Error.WriteLine("J = ({0:f3}, {1:f3}, {2:f3})", jx / jtot, jy / jtot, jz / jtot);
        s.AlignWithVector(jx, jy, jz);
        
        if (opts.Alignment)
        {
            double cos = Bullock.CalculateAlignment(ref s);
            Console.WriteLine("cos(theta_1/2) = {0:f4}", cos);
        }
        else if (opts.Bin)
        {
            //DataFile.Data bin = DoSphericalBins(s, opts.NumPerBin);
            DataFile.Table bin = Bullock.CalculateShells(s, opts.NumPerBin);
            bin.Print(opts.OutFile, new DataFile.Transform(Math.Log10));
        }
        else if (opts.Squared)
        {
            //DataFile.Data spec_j = CalculateSpecificJ(ref s, opts.NumPerBin, opts.Relative);
            DataFile.Table spec_j = Bullock.CalculateAngularMomentumDistributionJ2(ref s, opts.NumPerBin);
            // now fix error bars
            for (int i = 0; i < spec_j.Length; i++)
            {
                spec_j[i, 0] = Math.Log10(spec_j[i, 0]);
                spec_j[i, 1] = Math.Log10(spec_j[i, 1]);
                spec_j[i, 2] = Math.Log10(spec_j[i, 2]);
                spec_j[i, 3] = Math.Log10(spec_j[i, 3]);
                spec_j[i, 2] = 1.0;//spec_j[i, 0] - spec_j[i, 2];
                spec_j[i, 3] = 1.0;//spec_j[i, 3] - spec_j[i, 0];
            }
            spec_j.Print(opts.OutFile);
        }
        else
        {   
            //DataFile.Data spec_j = CalculateSpecificJ(ref s, opts.NumPerBin, opts.Relative);
            DataFile.Table spec_j = Bullock.CalculateAngularMomentumDistribution(ref s, opts.NumPerBin);
            // now fix error bars
            for (int i = 0; i < spec_j.Length; i++)
            {
                spec_j[i, 0] = Math.Log10(spec_j[i, 0]);
                spec_j[i, 1] = Math.Log10(spec_j[i, 1]);
                spec_j[i, 2] = Math.Log10(spec_j[i, 2]);
                spec_j[i, 3] = Math.Log10(spec_j[i, 3]);
                spec_j[i, 2] = 0.0; //spec_j[i, 0] - spec_j[i, 2];
                spec_j[i, 3] = 0.0; //spec_j[i, 3] - spec_j[i, 0];
            }
            spec_j.Print(opts.OutFile);
        }
    }
    
    static void CalculateTotalJVector(NBodySystem s, out double jx, out double jy, out double jz)
    {
        jx = jy = jz = 0.0;
        foreach (Particle p in s)
        {
            jx += p.Mass * (p.Y * p.Vz - p.Z * p.Vy);
            jy += p.Mass * (p.Z * p.Vx - p.X * p.Vz);
            jz += p.Mass * (p.X * p.Vy - p.Y * p.Vx);
        }
    }
    
    static void CalculateAlignment(ref NBodySystem s)
    {
        s.Sort();
        double m_tot = s.Mass;
        double m_half = m_tot / 2.0;
        
        double m_check = 0.0;
        int half = 0;
        while (m_check < m_half)
            m_check += s[half++].Mass;
        
        double jx1, jy1, jz1;
        jx1 = jy1 = jz1 = 0.0;
        for (int i = 0; i < half; i++)
        {
            jx1 += s[i].Mass * (s[i].Y * s[i].Vz - s[i].Z * s[i].Vy);
            jy1 += s[i].Mass * (s[i].Z * s[i].Vx - s[i].X * s[i].Vz);
            jz1 += s[i].Mass * (s[i].X * s[i].Vy - s[i].Y * s[i].Vx);
        }
        double abs1 = Math.Sqrt(jx1 * jx1 + jy1 * jy1 + jz1 * jz1);
        jx1 /= abs1;
        jy1 /= abs1;
        jz1 /= abs1;
        Console.WriteLine("Inner halo (r < {0:f4}):", s[half].Radius);
        Console.WriteLine("  jx = {0:f4}\n  jy = {1:f4}\n  jz = {2:f4}", jx1, jy1, jz1);
        
        double jx2, jy2, jz2;
        jx2 = jy2 = jz2 = 0.0;
        for (int i = half; i < s.NumParts; i++)
        {
            jx2 += s[i].Mass * (s[i].Y * s[i].Vz - s[i].Z * s[i].Vy);
            jy2 += s[i].Mass * (s[i].Z * s[i].Vx - s[i].X * s[i].Vz);
            jz2 += s[i].Mass * (s[i].X * s[i].Vy - s[i].Y * s[i].Vx);
        }
        double abs2 = Math.Sqrt(jx2 * jx2 + jy2 * jy2 + jz2 * jz2);
        jx2 /= abs2;
        jy2 /= abs2;
        jz2 /= abs2;
        Console.WriteLine("Out halo (r > {0:f4}):", s[half].Radius);
        Console.WriteLine("  jx = {0:f4}\n  jy = {1:f4}\n  jz = {2:f4}", jx2, jy2, jz2);
        
        double cos = (jx1 * jx2 + jy1 * jy2 + jz1 * jz2);
        
        Console.WriteLine("cos(theta_1/2) = {0:f4}", cos);
    }
    
    static DataFile.Table DoSphericalBins(NBodySystem s, int num_per_bin)
    {
        int num_bins = (int)(Math.Ceiling((double)s.NumParts / (double)num_per_bin));
        double rmax = s.Rmax + 1e-10;
        double rmin = s.Rmin;
        double deltaR = Math.Log10(rmax / rmin) / (double)num_bins;
        
        int[] num = new int[num_bins];
        double[] jxbin = new double[num_bins];
        double[] jybin = new double[num_bins];
        double[] jzbin = new double[num_bins];
        
        DataFile.Table data = new DataFile.Table(num_bins, 3, "Angular momentum in spherical shells");
        for (int i = 0; i < num_bins; i++)
        {
            jxbin[i] = jybin[i] = jzbin[i] = 0.0;
            data[i, 0] = 0.0; // radius
            data[i, 1] = 0.0; // enclosed mass
            data[i, 2] = 0.0; // average angular momentum in shell
            num[i] = 0;
        }
        
        for (int i = 0; i < s.NumParts; i++)
        {
            int pos = (int)((Math.Log10(s[i].Radius) - Math.Log10(rmin)) / deltaR);
            jxbin[pos] += s[i].Mass * s[i].Jx;
            jybin[pos] += s[i].Mass * s[i].Jy;
            jzbin[pos] += s[i].Mass * s[i].Jz;
            num[pos]++;
        }
        
        double m_enc = 0.0;
        double jmax = 0.0;
        for (int i = 0; i < num_bins; i++)
        {
            m_enc += (double)num[i] * s[0].Mass;
            data[i, 1] = m_enc;
            data[i, 0] = Math.Pow(10.0, Math.Log10(rmin) + 0.5 * deltaR + deltaR * i);
            if (num[i] > 0)
                data[i, 2] = Math.Sqrt(jxbin[i] * jxbin[i] + jybin[i] * jybin[i] + jzbin[i] * jzbin[i]);
            else
                data[i, 2] = 0.0;
            if (data[i, 2] > jmax) jmax = data[i, 2];
        }
        
        // normalize data ?
        for (int i = 0; i < num_bins; i++)
        {
            data[i, 0] /= rmax;
            data[i, 1] /= m_enc;
            data[i, 2] /= jmax;
        }
        
        return data;
    }   
    
    static DataFile.Table CalculateSpecificJ(ref NBodySystem s, int num_per_bin, bool relative)
    {
        int num_bins = (int)(Math.Ceiling((double)s.NumParts / (double)num_per_bin));
        double[] jcells = new double[num_bins * 3];
        double[] mcells = new double[num_bins * 3];
        s.Sort();
        if (use_squared)
            split_into_cells_j2(s, jcells, mcells, num_per_bin);
        else
            split_into_cells_jz(s, jcells, mcells, num_per_bin);
        
        IComparer comp = new ReverseSorter();
        Array.Sort(jcells, mcells, comp);
        
        DataFile.Table data = new DataFile.Table(jcells.Length, 2, "Specific angular momentum");
        double jmax = jcells[0];
        double mmax = 0.0;
        for (int i = 0; i < jcells.Length; i++)
        {
            data[i, 0] = jcells[i];
            for (int j = i + 1; j < jcells.Length; j++)
            {
                if (jcells[j] > 0)  // only include if positive
                {
                    data[i, 1] += mcells[j];
                }
            }
            if (jcells[i] > jmax) jmax = jcells[i];
            if (jcells[i] > 0.0) mmax += mcells[i];
        }
        
        Console.Error.WriteLine("Jmax = {0}", jmax);
        Console.Error.WriteLine("Mmax = {0}", mmax);
        Console.Error.WriteLine("Mtot = {0}", s.Mass);
        
        // now make the data relative
        if (relative)
        {
            double Mtot = s.Mass;
            for (int i = 0; i < data.Length; i++)
            {
                data[i, 0] /= jmax;
                data[i, 1] /= mmax;
            }
        }
        
        return data;
    }
    
    static void split_into_cells_jz(NBodySystem s, double[] jcells, double[] mcells, int n)
    {
        double vol1 = Math.Sqrt(5.0) / 3.0;
        double vol2 = Math.Sqrt(8.0) / 3.0;
        // system is sorted, so just put first n into bin 1, etc
        int c = 0;
        for (int i = 0; i < s.NumParts; i++)
        {
            double jz = s[i].Mass * s[i].Jz;
            double sin = Math.Sin(s[i].Theta);
            if (sin < vol1)
            {
                jcells[c + 0] += jz;
                mcells[c + 0] += s[i].Mass;
            }
            else if (sin < vol2)
            {
                jcells[c + 1] += jz;
                mcells[c + 1] += s[i].Mass;
            }
            else
            {
                jcells[c + 2] += jz;
                mcells[c + 2] += s[i].Mass;
            }
            
            if ((i+1) % n == 0)
                c += 3;
        }
        
        // how many cells have negative jz?
        int c_neg = 0;
        foreach (double x in jcells)
        {
            if (x < 0)
                c_neg++;
        }
        double frac_neg = (double)c_neg / (double)jcells.Length;
        Console.WriteLine("{0:f2}% of cells are negatively aligned", 100.0 * frac_neg);
    }
    
    static void split_into_cells_j2(NBodySystem s, double[] jcells, double[] mcells, int n)
    {
        double vol1 = Math.Sqrt(5.0) / 3.0;
        double vol2 = Math.Sqrt(8.0) / 3.0;
        // system is sorted, so just put first n into bin 1, etc
        int c = 0;
        double[] jxcells = new double[jcells.Length];
        double[] jycells = new double[jcells.Length];
        double[] jzcells = new double[jcells.Length];
        for (int i = 0; i < s.NumParts; i++)
        {
            double jx = s[i].Mass * s[i].Jx;
            double jy = s[i].Mass * s[i].Jy;
            double jz = s[i].Mass * s[i].Jz;
            double sin = Math.Sin(s[i].Theta);
            if (sin < vol1)
            {
                jxcells[c + 0] += jx;
                jycells[c + 0] += jy;
                jzcells[c + 0] += jz;
                mcells[c + 0] += s[i].Mass;
            }
            else if (sin < vol2)
            {
                jxcells[c + 1] += jx;
                jycells[c + 1] += jy;
                jzcells[c + 1] += jz;
                mcells[c + 1] += s[i].Mass;
            }
            else
            {
                jxcells[c + 2] += jx;
                jycells[c + 2] += jy;
                jzcells[c + 2] += jz;
                mcells[c + 2] += s[i].Mass;
            }
            
            if ((i+1) % n == 0)
                c += 3;
        }
        
        for (int i = 0; i < jcells.Length; i++)
        {
            jcells[i] = Math.Sqrt(jxcells[i] * jxcells[i] +
                                  jycells[i] * jycells[i] +
                                  jzcells[i] * jzcells[i]);
        }
    }
    
    public class ReverseSorter : IComparer 
    {
        public int Compare( Object x, Object y )
        {
            double xx = (double)x;
            double yy = (double)y;
            if (xx < yy)
                return 1;
            else if (yy < xx)
                return -1;
            else
                return 0;
        }
        
        int Compare( double x, double y )
        {
            if (x < y)
                return 1;
            else if (y < x)
                return -1;
            else
                return 0;
        }
    }
}
